import os
from peft import PeftModel, LoraConfig, TaskType
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载预训练模型和tokenizer
model_id = "./output2/Qwen2.5-VL-7B/"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

# 定义Lora配置
val_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    inference_mode=True,  # 推理模式
    r=64,  # Lora 秩
    lora_alpha=16,  # Lora alpha，具体作用参见 Lora 原理
    lora_dropout=0.05,  # Dropout 比例
    bias="none",
)

# 应用Lora配置并加载微调后的模型
val_peft_model = PeftModel.from_pretrained(model=model, model_id=model_id, config=val_config)

# 合并Lora权重到原始模型中
val_peft_model.merge_and_unload()

# 保存合并后的模型
merged_model_path = "./merged_model"
model.save_pretrained(merged_model_path)
tokenizer.save_pretrained(merged_model_path)

print(f"合并后的模型已保存至 {merged_model_path}")